import argparse
import os
import pickle
import warnings
import xml.etree.ElementTree as ET

import cv2
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import scipy.io as sio
import spectral.io.envi as envi
import torch
import torch.distributed as dist
import torch.nn as nn
from numpy import flip
from sklearn.metrics import accuracy_score
from spectral import imshow, view_cube
from torch.nn.parallel import DistributedDataParallel as DDP
from tqdm import tqdm
# from ViT import ViT
from vit_pytorch import ViT

from AutoGPU import autoGPU
from GAN_training_utils import TrainProcess, setup_seed
# from AutoGPU import autoGPU
from models import (_1DCNN, _2DCNN, _3DCNN, _3DCNN_1DCNN, _3DCNN_AM, PURE3DCNN,
                    PURE3DCNN_2AM, SAE, SAE_AM, DBDA_network, HamidaEtAl,
                    LeeEtAl, SSRN_network, myknn, mysvm)
from myTrans2 import Generator
from NViT import myViT
from utils import DataPreProcess, myplot, plot, setpath, splitdata

parser = argparse.ArgumentParser()
parser.add_argument('--dataset', type=str,  default='./pathology/data/032370b-20x-roi2', metavar='D',
                        help='the dataset path you load')
parser.add_argument('--trial_number', type=int,  default=1, metavar='T',
                    help='the time you do this trial')
parser.add_argument('--train_number', type=int,  default=0.05, metavar='NTr', 
                        help='number of training set')
parser.add_argument('--valid_number', type=int,  default=0.05, metavar='NVa',
                        help='number of valid set')
parser.add_argument('--test_number', type=int,  default=0.9, metavar='NTe',
                        help='number of test set')
parser.add_argument('--patchsize', type=int,  default=9, metavar='P',
                    help='patchsize of data') 
parser.add_argument('--modelname', type=str,  default='PURE3DCNN', metavar='P', help='which model to choose') 
# parser.add_argument('--gpu_ids', type=int,  default=-1, metavar='G',
#                     help='which gpu to use')

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 
args = parser.parse_args()

dataset = args.dataset
NTr = args.train_number
trialnumber = args.trial_number
NTe = args.test_number
NVa = args.valid_number
patchsize = args.patchsize
modelname = args.modelname

dark_ref = envi.open('%s.hdr'%dataset, '%s.raw'%dataset)
dark_nparr = np.array(dark_ref.load())
dark_nparr = flip(dark_nparr, 0)
tree = ET.parse('%s.xml'%dataset)
plt.figure()
plt.imshow(dark_nparr[:,:,10])
plt.savefig('test.jpg') 
    # get root element
polygon = tree.getroot().findall('object/polygon')

mask = np.zeros((dark_nparr.shape[0], dark_nparr.shape[1]), dtype="uint8")

IMAGE = np.flipud(dark_nparr)
GND = mask  


resultpath, imagepath, datapath = setpath(dataset, trialnumber , NTr, NVa, NTe, modelname)

with open(resultpath + 'result.pkl', 'rb') as f:
    result = pickle.load(f)

print(result.conf_mat.trace()/result.conf_mat.sum())

print(accuracy_score(result.y_true, result.y_pre))
resultpath, imagepath, datapath = setpath(dataset, trialnumber , NTr, NVa, NTe, modelname)
spliteddata = splitdata(IMAGE, GND, datapath , trainnum=NTr, validnum=NVa, testnum=NTe)
trainpos, traingt = DataPreProcess.parsespdata(spliteddata['train'])
testpos, testgt = DataPreProcess.parsespdata(spliteddata['test'])
processeddata = DataPreProcess(IMAGE, patchsize, datapath, 1).processeddata
imgPos = np.array(trainpos + testpos)
imgGt = np.array(traingt + result.y_pre)

myplot(processeddata, IMAGE, imagepath, result)

img = np.zeros(IMAGE.shape[:2]+(1,), dtype=int)
for p_idx, p in enumerate(imgPos):
    img[tuple(p)] = imgGt[p_idx]


patchradius = 10 


for p_idx, p in enumerate(tqdm(testpos)):

    if p[0]>patchradius and p[0]<IMAGE.shape[0]-patchradius and \
        p[1]>patchradius and p[1]<IMAGE.shape[1]-patchradius:
        sum = 0
        for i in range(p[0]-patchradius, p[0]+patchradius+1):     
            for j in range(p[1]-patchradius, p[1]+patchradius+1):                               
                sum += img[i, j]
        if sum >= (2*patchradius+1)**2*0.8:
            result.y_pre[p_idx] = 1
            img[p[0], p[1]] = 1
        elif sum <= (2*patchradius+1)**2*0.2:
            result.y_pre[p_idx] = 0
            img[p[0], p[1]] = 0
            

imagepath = imagepath + '/filter/'
if not os.path.exists(imagepath):
    os.makedirs(imagepath)
print(accuracy_score(result.y_true, result.y_pre))
myplot(processeddata, IMAGE, imagepath, result)
print('end')
